# encoding: utf-8
"""
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    created by lane.chang on '09/09/2024'
    comment: basic model
"""
import asyncio
import numpy as np
import ujson
from pydantic import Field
from laner.pydantic import BaseModel
from sqlalchemy.util import asyncio
from datetime import datetime

from project.core.base import UEnum, Element
from project.lib.common import get_redis
from project.model.chatgpt import ChatGpt

CONTEXT_TOKEN_LIMIT = 3000


class Role(UEnum):
    """ 角色
    """
    UDI = Element('udi', '酒店管理人员', bot_id='udi', user_id='uu01')
    CAREANA = Element('careana', '医疗保健方面的专家', bot_id='careana', user_id='nn01')
    HELPER = Element('helper', '助手', bot_id='helper', user_id='hh01')


class Question(BaseModel):
    """ 固定问题
    """
    index: str = Field('', title='问题')
    content: str = Field('', title='答案')

    def to_qa(self):
        """
        :return:
        """
        return f'Q: {self.index}\nA: {self.content}\n'


class Knowledge(BaseModel):
    """ 知识库
    """
    key: str = Field('', title='知识库key')
    questions: list[Question] = Field([], title='知识库')
    emb_model_name: str = Field('text-embedding-ada-002', title='')
    sources: list = Field([])
    embeddings: list = Field([])

    def __init__(self, **kwargs):
        """
        """
        super(Knowledge, self).__init__(**kwargs)

        key = self.get_rds_key()
        rds = get_redis()

        self.sources = []
        self.embeddings = []
        if rds.exists(key):
            result = rds.get(key)
            result = ujson.loads(result)

            self.sources = result['sources']
            self.embeddings = result['embeddings']

    def get_rds_key(self):
        """
        :return:
        """
        return f'knowledge_{self.key}'

    async def generate_vector_db(self):
        """ 生成向量数据
        :return:
        """
        embeddings = []
        sources = []
        chat_gpt = ChatGpt(emb_model_name=self.emb_model_name)
        for question in self.questions:
            sources.append({'question': question.index, 'answer': question.content})

            text = f'{question.index} {question.content}'
            result = await chat_gpt.get_embedding(text)
            embeddings.append(result)

        return {'sources': sources, 'embeddings': embeddings}

    async def flush_vector_db(self):
        """ 保存向量数据
        :return:
        """
        key = self.get_rds_key()
        rds = get_redis()

        result = await self.generate_vector_db()
        rds.set(key, ujson.dumps(result, ensure_ascii=False))

    async def get_info(self):
        """
        :return:
        """
        key = self.get_rds_key()
        rds = get_redis()
        if not rds.exists(key):
            return {}

        result = rds.get(key)
        return ujson.loads(result)

    async def order_document_sections_by_query_similarity(self, text):
        """
        :param text:
        :return:
        """
        chat_gpt = ChatGpt(emb_model_name=self.emb_model_name)
        text_embedding = await chat_gpt.get_embedding(text)

        document_similarities = sorted([
            (np.dot(text_embedding, doc_embedding), doc_index) for doc_index, doc_embedding in
            enumerate(self.embeddings)
        ], reverse=True, key=lambda x: x[0])

        return document_similarities

    async def fetch_relevant(self, text):
        """
        :param text:
        :return:
        """
        ordered_candidates = await self.order_document_sections_by_query_similarity(text)
        ctx = ""
        ctx_lst = []
        for idx, candi in enumerate(ordered_candidates):
            next = ctx + " " + self.sources[candi[1]]['question'] + self.sources[candi[1]]['answer']
            if len(next) > CONTEXT_TOKEN_LIMIT:  # 最大的token字数
                break
            ctx = next
            ctx_lst.append(self.sources[candi[1]])
        return ctx, ctx_lst

    async def retrieve(self, text):
        """
        :param text:
        :return:
        """
        if not self.sources or not self.embeddings:
            return []

        ctx, ctx_lst = await self.fetch_relevant(text)
        # 最多取3条现在
        ctx_lst = ctx_lst[:3]

        return ctx_lst


class User(BaseModel):
    """ 用户信息
    """
    bot_id: str = Field('udi', title='数字人id', description='udi/careana/ 默认udi')
    user_id: str = Field('uu01', title='用户id', description='uu01/nn01 默认uu01')
    dev_id: str = Field('dev01', title='设备id', description='设备id')
    role: str = Field('客人', title='用户角色')


class Room(BaseModel):
    """ 房间信息
    """
    room_type: list[str] = Field('会议室', title='房间类型')
    building_name: str = Field('', title='楼栋名称')
    floor_name: str = Field('', title='楼层名称')
    room_name: str = Field('', title='房间名称')
    area_location: str = Field('', title='区域位置')
    head_count: str = Field(0, title='人数')

    @property
    def agent_dict(self):
        """
        :return:
        """
        return {
            "房间类型": ",".join(self.room_type),
            "楼栋名称": self.building_name,
            "楼层名称": self.floor_name,
            "房间名称": self.room_name,
            "区域位置": f'{self.building_name}{self.floor_name}{self.room_name}',
            "人数": self.head_count
        }

class KlItem(BaseModel):
    """ 知识库单项信息
    """
    question: str = Field('', title='问题')
    answer: str = Field('', title='答案')


class ChatMessage(BaseModel):
    """ 聊天信息
    """
    role: str = Field('', title='角色', description='角色 user(用户)/agent(智能体)')
    message: str = Field('', title='聊天文本')
    create_time: str = Field('', title='聊天时间')

    def __str__(self):
        """
        :return:
        """
        ret_data = ''
        chat_seq = datetime.strptime(self.create_time, '%Y-%m-%d %H:%M:%S').strftime('%Y%m%d%H%M%S')
        if self.role == 'user':
            ret_data = '对话序号:' + str(chat_seq) + f"\t\t用户对你说:" + self.message
        else:
            ret_data = '对话序号:' + str(chat_seq) + f"\t\t你对用户说:" + self.message

        return ret_data


if __name__ == '__main__':

    pass
    # knowledge = Knowledge(
    #     key='su_kang_yang',
    #     questions=[
    #         Question(index='酒店的室内恒温泳池开放时间是什么时候？', content='我们的室内恒温泳池每周七天开放，从早上6:30至晚上11:00，确保您在任何时间都能享受到游泳的乐趣。'),
    #         Question(index='酒店提供哪些餐饮服务？', content='酒店设有盛宴标帜餐厅、大堂吧和采悦轩中餐厅，提供零点服务、早餐自助、午餐和晚餐自助。盛宴标帜餐厅位于2楼，大堂吧位于1楼，采悦轩中餐厅提供包房商务套餐和特色菜品套餐。')
    #     ]
    # )

    knowledge = Knowledge(key=15)
    result = asyncio.run(knowledge.retrieve('酒店都有哪些餐饮服务?'))
    print(result)



